package com.zzl.study.auth.security.handler;

import com.fasterxml.jackson.core.type.TypeReference;
import com.zzl.study.auth.domain.Response;
import com.zzl.study.auth.utils.JacksonUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.core.Authentication;
import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
import org.springframework.security.crypto.keygen.StringKeyGenerator;
import org.springframework.security.oauth2.client.ReactiveOAuth2AuthorizedClientService;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.user.DefaultOidcUser;
import org.springframework.security.web.server.WebFilterExchange;
import org.springframework.security.web.server.authentication.logout.ServerLogoutHandler;
import org.springframework.security.web.server.context.ServerSecurityContextRepository;
import org.springframework.util.MultiValueMap;
import org.springframework.util.ObjectUtils;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;

import java.util.Base64;
import java.util.Map;

/**
 * @author: zhangzl
 * @date: 2024/6/14 11:03
 * @version: 1.0
 * @description: 自定义退出登录处理器
 */
public class GatewayLogoutHandler implements ServerLogoutHandler {

    private static final String FAIL_LOGOUT_RESPONSE_ERROR_CODE = "fail_logout";

    private static final StringKeyGenerator DEFAULT_STATE_GENERATOR = new Base64StringKeyGenerator(Base64.getUrlEncoder());

    private final ReactiveOAuth2AuthorizedClientService reactiveOAuth2AuthorizedClientService;

    private final ServerSecurityContextRepository serverSecurityContextRepository;

    private final WebClient webClient = WebClient.builder().build();

    public GatewayLogoutHandler(ReactiveOAuth2AuthorizedClientService reactiveOAuth2AuthorizedClientService, ServerSecurityContextRepository serverSecurityContextRepository) {
        this.reactiveOAuth2AuthorizedClientService = reactiveOAuth2AuthorizedClientService;
        this.serverSecurityContextRepository = serverSecurityContextRepository;
    }

    @Override
    public Mono<Void> logout(WebFilterExchange exchange, Authentication authentication) {
        OAuth2AuthenticationToken token = (OAuth2AuthenticationToken) authentication;
        String authorizedClientRegistrationId = token.getAuthorizedClientRegistrationId();
        String name = token.getName();
        DefaultOidcUser principal = (DefaultOidcUser) token.getPrincipal();
        String idTokenHint = principal.getIdToken().getTokenValue();
        return this.reactiveOAuth2AuthorizedClientService.loadAuthorizedClient(authorizedClientRegistrationId, name)
                .flatMap(authorizedClient -> {
                    String tokenValue = authorizedClient.getAccessToken().getTokenValue();
                    ClientRegistration clientRegistration = authorizedClient.getClientRegistration();
                    Map<String, Object> configurationMetadata = clientRegistration.getProviderDetails().getConfigurationMetadata();
                    String logoutUri = String.valueOf(configurationMetadata.get("end_session_endpoint"));
                    String url = UriComponentsBuilder.fromUriString(logoutUri)
                            .queryParam("id_token_hint", idTokenHint)
                            .queryParam(OAuth2ParameterNames.CLIENT_ID, clientRegistration.getClientId())
//                            .queryParam("post_logout_redirect_uri", "http://micro-service-gateway/login-ui")
//                            .queryParam(OAuth2ParameterNames.STATE, DEFAULT_STATE_GENERATOR.generateKey())
                            .build()
                            .toString();
                    String nonceId = getNonceId(exchange);
                    return this.webClient.get()
                            .uri(url)
                            .header(HttpHeaders.ACCEPT, MediaType.APPLICATION_JSON_VALUE)
                            .header("nonceId", nonceId)
                            .headers((headers) -> headers.setBearerAuth(tokenValue))
                            .retrieve()
                            .onStatus(HttpStatusCode::isError, (response) ->
                                    response.bodyToMono(String.class).flatMap(body -> {
                                        Response<String> restResponse = JacksonUtils.toObject(body, new TypeReference<Response<String>>() {
                                        });
                                        OAuth2Error oauth2Error = new OAuth2Error(FAIL_LOGOUT_RESPONSE_ERROR_CODE, restResponse.getMsg(), null);
                                        return Mono.error(new OAuth2AuthenticationException(oauth2Error, oauth2Error.toString()));
                                    })
                            )
                            .bodyToMono(String.class)
                            .flatMap(body -> {
                                Response<String> restResponse = JacksonUtils.toObject(body, new TypeReference<Response<String>>() {
                                });
                                if (restResponse.getCode() != 0) {
                                    OAuth2Error oauth2Error = new OAuth2Error(FAIL_LOGOUT_RESPONSE_ERROR_CODE, restResponse.getMsg(), null);
                                    return Mono.error(new OAuth2AuthenticationException(oauth2Error, oauth2Error.toString()));
                                }
                                return this.reactiveOAuth2AuthorizedClientService.removeAuthorizedClient(authorizedClientRegistrationId, name)
                                        .then(this.serverSecurityContextRepository.save(exchange.getExchange(), null));
                            });
                });

    }

    public String getNonceId(WebFilterExchange exchange) {
        ServerHttpRequest request = exchange.getExchange().getRequest();
        HttpHeaders headers = request.getHeaders();
        MultiValueMap<String, String> queryParams = request.getQueryParams();
        String nonce = headers.getFirst("nonceId");
        if (ObjectUtils.isEmpty(nonce)) {
            nonce = queryParams.getFirst("nonceId");
        }
        return nonce;
    }

}
